from typing import Optional, List
from asyncio import Queue
import os
import asyncio # 导入 asyncio

# DGA 数据库文件路径
DGA_DATABASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), 'docs', '2025-04-05')

# 全局变量，用于存储加载的 DGA 域名集合
DGA_DOMAINS = set()

# 在模块加载时读取 DGA 域名数据库
try:
    with open(DGA_DATABASE_PATH, 'r', encoding='utf-8') as f:
        for line in f:
            domain = line.strip().lower() # 读取并转换为小写
            if domain:
                DGA_DOMAINS.add(domain)
    print(f"[DGA Detector] 成功加载 {len(DGA_DOMAINS)} 个 DGA 域名。")
except FileNotFoundError:
    print(f"[DGA Detector] 警告: 未找到 DGA 数据库文件: {DGA_DATABASE_PATH}。DGA 检测功能可能受限。")
except Exception as e:
    print(f"[DGA Detector] 警告: 加载 DGA 数据库文件时发生错误: {e}。DGA 检测功能可能受限。")

async def _send_status(status_queue: Optional[Queue], tool_name: str, message_type: str, data: dict):
    """辅助函数：向状态队列发送消息。"""
    if status_queue:
        await status_queue.put({
            "type": message_type,
            "step": tool_name,
            "data": data
        })

async def detect_dga_domain(domain_name: str, status_queue: Optional[Queue] = None) -> dict:
    """
    检测一个域名是否是已知的 DGA 域名。
    """
    tool_name = "DGA_Detector"
    await _send_status(status_queue, tool_name, "tool_start", {"domain_name": domain_name})

    is_dga = False
    reason = "域名未在已知DGA数据库中发现。"

    if not domain_name:
        reason = "域名为空。"
    elif domain_name.lower() in DGA_DOMAINS:
        is_dga = True
        reason = "域名在已知DGA数据库中发现。"

    analysis_result = {
        "is_dga": is_dga,
        "reason": reason
    }

    await _send_status(status_queue, tool_name, "tool_end", {"domain_name": domain_name, "result": analysis_result})
    return {"source": "Internal-DGA-Detector", "data": analysis_result, "status": "success"}

async def batch_detect_dga_domains(domain_list: List[str], status_queue: Optional[Queue] = None) -> List[dict]:
    """
    批量检测域名是否是已知的 DGA 域名。
    """
    tool_name = "Batch_DGA_Detector"
    await _send_status(status_queue, tool_name, "tool_start", {"domain_count": len(domain_list)}) 

    results = []
    tasks = []
    for domain in domain_list:
        # 为每个域名创建一个异步任务
        tasks.append(detect_dga_domain(domain, status_queue))
    
    # 并发执行所有检测任务
    domain_results = await asyncio.gather(*tasks)

    for domain, result in zip(domain_list, domain_results):
        results.append({
            "domain": domain,
            "is_dga": result["data"]["is_dga"],
            "reason": result["data"]["reason"],
            "source": result["source"]
        })

    await _send_status(status_queue, tool_name, "tool_end", {"results_count": len(results)})
    return results